package org.bkhech.mybatis.extend.generator.plugins;

import org.bkhech.mybatis.extend.generator.constant.GeneratorConstant;
import org.bkhech.mybatis.extend.generator.util.GeneratorUtil;
import org.mybatis.generator.api.IntrospectedColumn;
import org.mybatis.generator.api.IntrospectedTable;
import org.mybatis.generator.api.PluginAdapter;
import org.mybatis.generator.api.dom.java.Field;
import org.mybatis.generator.api.dom.java.FullyQualifiedJavaType;
import org.mybatis.generator.api.dom.java.Method;
import org.mybatis.generator.api.dom.java.TopLevelClass;
import org.mybatis.generator.internal.util.StringUtility;

import java.util.HashSet;
import java.util.Iterator;
import java.util.List;

/**
 * 修改生成的Model，让其继承BaseModel
 *
 * Created by guowm on 2017/2/13.
 *
 * Updated
 *      支持表可以没有主键, 支持联合主键
 */
public class BaseModelGeneratorPlugin extends PluginAdapter {

    private static final String DEFAULT_PRIMARYKEY = "id";

    @Override
    public boolean validate(List<String> list) {
        String modelTargetPackage = this.properties.getProperty("modelTargetPackage");
        return StringUtility.stringHasValue(modelTargetPackage);
    }

    @Override
    public boolean modelBaseRecordClassGenerated(TopLevelClass topLevelClass, IntrospectedTable introspectedTable) {
        List<IntrospectedColumn> primaryKeyColumns = introspectedTable.getPrimaryKeyColumns();

        if (primaryKeyColumns.isEmpty()) {
            FullyQualifiedJavaType superClazzType = new FullyQualifiedJavaType(GeneratorConstant.BASE_MODEL_CLASS_PATH);
            topLevelClass.addImportedType(superClazzType);
            topLevelClass.setSuperClass(superClazzType);
        } else if (primaryKeyColumns.size() == 1) {
            FullyQualifiedJavaType pkType = primaryKeyColumns.get(0).getFullyQualifiedJavaType();

            FullyQualifiedJavaType superClazzType = new FullyQualifiedJavaType(GeneratorConstant.BASE_PRIMARY_KEY_MODEL_CLASS_PATH);
            topLevelClass.addImportedType(superClazzType);

            superClazzType.addTypeArgument(pkType);
            topLevelClass.setSuperClass(superClazzType);
            clearPrimaryKeyColums(topLevelClass, primaryKeyColumns);
        }

        topLevelClass.addJavaDocLine("/**");
        topLevelClass.addJavaDocLine(" * Model: " + topLevelClass.getType().getShortName());
        topLevelClass.addJavaDocLine(" * Table: " + introspectedTable.getFullyQualifiedTable().getIntrospectedTableName());
        topLevelClass.addJavaDocLine(" * Alias: " + introspectedTable.getFullyQualifiedTable().getAlias());
        topLevelClass.addJavaDocLine(" *");
        topLevelClass.addJavaDocLine(" * This Model generated by MyBatis Generator Extend.");
        topLevelClass.addJavaDocLine(" */");
        return super.modelBaseRecordClassGenerated(topLevelClass, introspectedTable);
    }

    @Override
    public boolean modelPrimaryKeyClassGenerated(TopLevelClass topLevelClass, IntrospectedTable introspectedTable) {
        List<IntrospectedColumn> primaryKeyColumns = introspectedTable.getPrimaryKeyColumns();
        if (!primaryKeyColumns.isEmpty() && primaryKeyColumns.size() > 1) {
            FullyQualifiedJavaType pkType = topLevelClass.getType();
            FullyQualifiedJavaType superClazzType = new FullyQualifiedJavaType(GeneratorConstant.BASE_MODEL_CLASS_PATH);
            topLevelClass.addImportedType(superClazzType);

            superClazzType.addTypeArgument(pkType);
            topLevelClass.setSuperClass(superClazzType);
        }
        topLevelClass.addJavaDocLine("/**");
        topLevelClass.addJavaDocLine(" * Primary: " + topLevelClass.getType().getShortName());
        topLevelClass.addJavaDocLine(" * Table: " + introspectedTable.getFullyQualifiedTable().getIntrospectedTableName());
        topLevelClass.addJavaDocLine(" * Alias: " + introspectedTable.getFullyQualifiedTable().getAlias());
        topLevelClass.addJavaDocLine(" *");
        topLevelClass.addJavaDocLine(" * This PrimaryKey generated by MyBatis Generator Extend.");
        topLevelClass.addJavaDocLine(" */");
        return super.modelPrimaryKeyClassGenerated(topLevelClass, introspectedTable);
    }

    /**
     * 去除子类中主键属性
     *
     * @param topLevelClass
     * @param primaryKeyColumns
     */
    private void clearPrimaryKeyColums(TopLevelClass topLevelClass, List<IntrospectedColumn> primaryKeyColumns) {
        if (primaryKeyColumns.isEmpty()) {
            return;
        }

        List<Field> fields = topLevelClass.getFields();
        Iterator<Field> fieldsIterator = fields.iterator();
        HashSet<Field> pkFields = new HashSet<Field>();

        List<Method> methods = topLevelClass.getMethods();
        Iterator<Method> methodsIterator = methods.iterator();
        HashSet<Method> pkMethods = new HashSet<Method>();

        while (fieldsIterator.hasNext()) {
            Field field = fieldsIterator.next();
            for (IntrospectedColumn column : primaryKeyColumns) {
                String pkName = column.getJavaProperty();
                if (field.getName().equals(pkName) && DEFAULT_PRIMARYKEY.equals(pkName)) {
                    pkFields.add(field);
                }
            }
        }
        fields.removeAll(pkFields);

        while (methodsIterator.hasNext()) {
            Method method = methodsIterator.next();
            for (IntrospectedColumn column : primaryKeyColumns) {
                String pkName = column.getJavaProperty();
                if (DEFAULT_PRIMARYKEY.equals(pkName)) {
                    String setter = "set" + GeneratorUtil.convertFieldName(pkName);
                    String getter = "get" + GeneratorUtil.convertFieldName(pkName);
                    if (method.getName().equals(setter) || method.getName().equals(getter)) {
                        pkMethods.add(method);
                    }
                }
            }
        }
        methods.removeAll(pkMethods);
    }
}
